#!/usr/bin/env python3
"""
向量化工具安装和设置脚本
"""

import os
import sys
import subprocess
import json
from pathlib import Path

def check_python_version():
    """检查Python版本"""
    if sys.version_info < (3, 8):
        print("错误: 需要Python 3.8或更高版本")
        return False
    print(f"✓ Python版本: {sys.version}")
    return True

def install_requirements():
    """安装依赖包"""
    requirements_file = "requirements_vectorize.txt"
    
    if not os.path.exists(requirements_file):
        print(f"错误: 找不到 {requirements_file}")
        return False
    
    try:
        print("正在安装依赖包...")
        subprocess.run([
            sys.executable, "-m", "pip", "install", "-r", requirements_file
        ], check=True)
        print("✓ 依赖包安装完成")
        return True
    except subprocess.CalledProcessError as e:
        print(f"错误: 安装依赖包失败 - {e}")
        return False

def check_ollama():
    """检查Ollama服务"""
    try:
        import requests
        response = requests.get("http://localhost:11434/api/tags", timeout=5)
        if response.status_code == 200:
            print("✓ Ollama服务运行正常")
            
            # 检查是否有mxbai-embed-large模型
            models = response.json().get("models", [])
            model_names = [model.get("name", "") for model in models]
            
            if any("mxbai-embed-large" in name for name in model_names):
                print("✓ mxbai-embed-large模型已安装")
            else:
                print("⚠ 警告: 未找到mxbai-embed-large模型")
                print("请运行: ollama pull mxbai-embed-large")
            
            return True
        else:
            print("⚠ 警告: Ollama服务响应异常")
            return False
    except Exception as e:
        print("⚠ 警告: 无法连接Ollama服务")
        print("请确保Ollama已安装并运行在 http://localhost:11434")
        print("安装Ollama: https://ollama.ai/")
        return False

def check_milvus():
    """检查Milvus连接"""
    try:
        from pymilvus import connections, utility
        connections.connect(host="localhost", port=19530)
        print("✓ Milvus服务连接成功")
        connections.disconnect("default")
        return True
    except Exception as e:
        print("⚠ 警告: 无法连接Milvus服务")
        print("请确保Milvus已安装并运行在 localhost:19530")
        print("Docker安装: docker run -p 19530:19530 -p 9091:9091 milvusdb/milvus:latest")
        return False

def create_sample_config():
    """创建示例配置文件"""
    config_file = "config.json"
    
    if os.path.exists(config_file):
        print(f"✓ 配置文件 {config_file} 已存在")
        return True
    
    config = {
        "embedding": {
            "model": "mxbai-embed-large",
            "url": "http://localhost:11434/api/embeddings",
            "dimension": 1024
        },
        "milvus": {
            "host": "localhost",
            "port": 19530,
            "collection_name": "md_documents"
        },
        "processing": {
            "chunk_size": 1000,
            "chunk_overlap": 200,
            "max_chunks_per_file": 1000
        },
        "input": {
            "md_directory": str(Path.cwd() / "md_output"),
            "file_patterns": ["*.md"]
        },
        "logging": {
            "level": "INFO",
            "log_file": "vectorization.log"
        }
    }
    
    try:
        with open(config_file, 'w', encoding='utf-8') as f:
            json.dump(config, f, indent=2, ensure_ascii=False)
        print(f"✓ 创建配置文件: {config_file}")
        return True
    except Exception as e:
        print(f"错误: 创建配置文件失败 - {e}")
        return False

def show_usage():
    """显示使用说明"""
    print("\n" + "="*60)
    print("向量化工具使用说明")
    print("="*60)
    print()
    print("1. 向量化MD文件:")
    print("   python run_vectorization.py")
    print()
    print("2. 搜索文档:")
    print("   python search_documents.py '搜索关键词'")
    print("   python search_documents.py -i  # 交互模式")
    print()
    print("3. 配置文件: config.json")
    print("   - 修改嵌入模型和Milvus连接设置")
    print("   - 调整文本分块参数")
    print()
    print("4. 日志文件: vectorization.log")
    print("   - 查看详细的处理日志")
    print()
    print("文件说明:")
    print("- vectorize_md_to_milvus.py: 核心向量化类")
    print("- run_vectorization.py: 执行向量化")
    print("- search_documents.py: 搜索工具")
    print("- config.json: 配置文件")
    print("- requirements_vectorize.txt: 依赖包")

def main():
    """主安装函数"""
    print("向量化工具安装和设置")
    print("="*40)
    
    success = True
    
    # 检查Python版本
    if not check_python_version():
        success = False
    
    # 安装依赖包
    if not install_requirements():
        success = False
    
    # 检查服务
    print("\n检查服务状态:")
    check_ollama()
    check_milvus()
    
    # 创建配置文件
    print("\n配置文件:")
    create_sample_config()
    
    if success:
        print("\n✓ 设置完成!")
        show_usage()
    else:
        print("\n⚠ 设置过程中遇到一些问题，请检查上述错误信息")
    
    return 0 if success else 1

if __name__ == "__main__":
    exit(main())
